__author__ = 'Vladimir Iglovikov'

import copy


class State:
    def __init__(self, inputParameter):
        if isinstance(inputParameter, list):
            self.nSites = len(inputParameter)
            self.up = [0] * self.nSites
            self.down = [0] * self.nSites
            for i in range(self.nSites):
                if inputParameter[i] == 1:
                    self.down[i] = 1
                elif inputParameter[i] == 2:
                    self.up[i] = 1
                elif inputParameter[i] == 3:
                    self.up[i] = 1
                    self.down[i] = 1
        elif isinstance(inputParameter, int):
            self.nSites = inputParameter
            self.up = [0] * self.nSites
            self.down = [0] * self.nSites




    """
    returns Sz of the state
    """
    def sz(self):
        return sum(self.up) - sum(self.down)

    """
    returns number of electrons
    """
    def nElectrons(self):
        return sum(self.up) + sum(self.down)

    """
    clones existing state
    """
    def clone(self):
        return copy.deepcopy(self)


    """
    annihilates spin from the i'th position with spin up
    """
    def lowerUp(self, i):
        assert self.up[i] == 1
        self.up[i] = 0

    """
    Creates spin on the i'th position with spin up
    """
    def raiseUp(self, i):
        assert self.up[i] == 0
        self.up[i] = 1

    """
    annihilates spin with a spin down at i'th position
    """
    def lowerDown(self, i):
        assert self.down[i] == 1
        self.down[i] = 0

    """
    creates spin with a spin down at the i'th position
    """
    def raiseDown(self, i):
        assert self.down[i] == 0
        self.down[i] == 1

    """
    returns integer corresponding to the binary number formed by the up spins
    """
    def binaryUp(self):
        return int(''.join([str(tx) for tx in self.up]), 2)

    """
    returns integer corresponding to the binary number formed by down spins
    """
    def binaryDown(self):
        return int(''.join([str(tx) for tx in self.down]), 2)

    def __str__(self):
        return """Down {down}
Up {up}""".format(up=self.up, down=self.down)
